## UC model for ex-ante real rate: RW + AR(1)

rm(list=ls())
graphics.off()
source("R/uc_fns.R")
source("R/data_fns.R")
library(KFAS)
library(dplyr)
set.seed(616)

## CHOOSE HERE

## (1) Short rate (1y real rate)
## cat("Estimation for 1y real rate:\n")
## data <- loadShortRate()
## filename <- "results/uc_estimates_y1.RData"

## (2) Long rate (10y real rate)
cat("Estimation for 10y real rate:\n")
data <- loadLongRate()
filename <- "results/uc_estimates_y10.RData"

##################################################

cat("Annual sample starts in", min(data$year), "and ends in", max(data$year), "\n")
cat("Filename for estimation results:", filename, "\n")
prior <- get_prior()

## MCMC
N <- 5     ## number of parallel chains/starting values MLE
M <- 20000  ## length of MCMC chain
T <- length(data$rr)
phi <- matrix(NA, N, M)
sigv2 <- matrix(NA, N, M)
sigeta2 <- matrix(NA, N, M)
x <- array(NA, c(N, M, 2, T))

start <- Sys.time()
for (i in 1:N) {
    ## initialize parameters
    pars <- getRandomPars()
    cat("Chain", i, "- initialized at:\n")
    cat("phi =", pars$phi, "sigv2 =", pars$sigv2, "sigeta2 =", pars$sigeta2, "\n")
    for (j in 1:M) {
        if (j %% 5000 == 0)
            cat("iteration", j, "\n")
        ## draw latent state variables
        mod <- ssm_uc(data$rr, pars, prior)
        sim <- simulateSSM(mod, type="states")  # simulation smoother
        pars$x <- t(sim[,,1])
        ## AR(p) parameters
        pars$phi <- draw_regression(pars$x[2,-1], pars$x[2,-T], pars$sigv2, 0, 1/prior$var_phi, restrict=TRUE)
        ## random walk innovation variance
        alpha1 <- prior$alpha_eta + (T-2)
        delta1 <- prior$delta_eta + sum(diff(pars$x[1,])^2)
        pars$sigeta2 <- 1/rgamma(1, alpha1/2, delta1/2)
        ## AR(p) innovation variance
        ## prior:
        G <- embed(pars$x[2,], 2); ydat <- G[,1]; xdat <- G[,-1,drop=FALSE]
        resid <- ydat - xdat * pars$phi
        alpha1 <- prior$alpha_v + (T-2)
        delta1 <- prior$delta_v + sum(resid^2)
        pars$sigv2 <- 1/rgamma(1, alpha1/2, delta1/2)
        ## save draws
        phi[i,j] <- pars$phi
        sigv2[i,j] <- pars$sigv2
        sigeta2[i,j] <- pars$sigeta2
        x[i,j,,] <- pars$x[1:2,]
    }
    cat("Duration so far:", format(Sys.time()-start), "\n")
}
cat(N*M, "iterations take\n")
print(Sys.time() - start)

## save full MCMC sample
mcmc_save <- list(phi=phi, sigv2=sigv2, sigeta2=sigeta2)
parnames <- c("phi", "sigv2", "sigeta2")
n <- length(parnames)

## flatten multiple chains into one chain and drop burn-in sample (first half)
ind <- rep((M/2+1):M, N) + M*rep(0:(N-1), each=M/2)
phi <- as.numeric(mcmc_save$phi)[ind]
sigv2 <- as.numeric(mcmc_save$sigv2)[ind]
sigeta2 <- as.numeric(mcmc_save$sigeta2)[ind]
X <- array(x, c(N*M, 2, T))[ind,,]
stopifnot(all(abs(phi)<1))
lambda <- sqrt(sigeta2/sigv2)*abs(1-phi)
print(summary(lambda))
## dev.new()
## plot(density(na.omit(lambda)))

## parameter estimates
tbl <- matrix(NA, n+1, 4)
rownames(tbl) <- c(parnames, "lambda (signal-to-noise)")
colnames(tbl) <- c("MCMC mean", "median", "LB", "UB")
## MCMC
theta <- cbind(phi, sigv2, sigeta2, lambda)
tbl[,1] <- colMeans(theta, na.rm=TRUE)   # posterior means
tbl[,2:4] <- t(apply(theta, 2, quantile, c(.5, .025, .975), na.rm=TRUE))
print(round(tbl,4))

## posterior distributions
sigu <- sqrt(sigeta2)
print("Posterior distribution of sigma_u:")
print(summary(sigu))
print("Posterior distribution of phi:")
print(summary(phi))
cat("# posterior correlations\n")
print(cor(theta))

## MCMC estimate of rstar
data$rstar.mean <- apply(X[,1,], 2, mean)
data$rstar.lb <- apply(X[,1,], 2, quantile, .025)
data$rstar.ub <- apply(X[,1,], 2, quantile, .975)

save(data, X, sigeta2, sigv2, phi, prior, file=filename)
